import sys
sys.path.append(r'/home/wangzheng/wangyangtest2/pyphot-engine/pyphot-engine')
import subprocess
import multiprocessing
import numpy as np
from numpy.fft import fftshift, fft, ifft, ifftshift
import time
import torch
from torch.fft import fftshift, fft, ifft, ifftshift
import math
from data import in_fiber,out_fiber
from multiprocessing import get_start_method, set_start_method,Process, Queue
import os

def process_polarization(data,
        sampling_rate,
        num_spans,
        beta2,
        delta_z,
        gamma,
        alpha,
        span_length,
        type_p,
        shared_dict,
        condition,
        output_queuetype,
        gpu_device):
    current_pid = os.getpid()
    print(f"当前进程PID:{current_pid},{shared_dict['flag']}")
    num_steps_per_span = int(span_length / delta_z)
    with torch.no_grad():
        length = np.size(data)
        process_type  = type_p
        torch.cuda.set_device(gpu_device)  # 选择GPU
        data = torch.from_numpy(data).to("cuda")
        data = data.to("cuda")
        data = data.clone().detach().to(dtype=torch.complex64)
        data_normal = data / torch.sqrt(torch.mean(abs(data) ** 2))

        power_dbm_set = -6
        power_mw = 10 ** (power_dbm_set / 10)
        power_w = power_mw / 1000
        power_origin = torch.mean(torch.abs(data_normal) ** 2)
        data = data_normal / torch.sqrt(power_origin / power_w)
        power_temp = torch.mean(torch.abs(data) ** 2)

        omega = np.arange(-sampling_rate / 2, sampling_rate / 2, sampling_rate / length)
        omega = np.reshape(omega, (-1, 1))
        omega = torch.from_numpy(omega)
        omega = omega.to("cuda")

        data_fft = fftshift(fft(data, dim=0), dim=0) / length  # 将数据变到频率并将零值移到原点  # 将数据变到频率并将零值移到原点

        for i in range(1, num_spans + 1):
            for j in range(1, num_steps_per_span + 1):
                # 1/2 信号衰减
                data_fft = data_fft * math.exp(-alpha * (delta_z / 2))

                # 1/2 光纤色散
                # 偏振
                data_fft = data_fft * torch.exp(-1j * (beta2 / 2) * (omega ** 2) * (delta_z / 2))

                # 光纤非线性
                # 偏振
                data_t = ifft(ifftshift(data_fft, dim=0), dim=0) * data_fft.size(dim=-1)
                '''在这里将data_t数据添加到共享内存中'''
                # 设置共享字典shared_dict["signal"]初始值
                if (shared_dict["signal"] is None and shared_dict["flag"] == "y"):
                    #pytorch张量不能当做共享内容用于进程间的使用
                    data_cpu = data_t.to("cpu").numpy()
                    shared_dict["signal"] = data_cpu
                    shared_dict["flag"] = "x"
                    print(f"执行初始化进程PID:{current_pid},{shared_dict['flag']}")

                with condition:
                    while shared_dict["flag"] != process_type:
                        condition.wait()  # 等待条件变量满足
                    # 当条件满足，继续执行下面的代码
                    print(f"修改后进程PID:{current_pid},{shared_dict['flag']}")
                    signal = shared_dict["signal"]
                    signal = torch.tensor(signal, dtype=torch.complex64).to("cuda")
                    # 在这里使用 signal 做一些操作
                    data_t = data_t * torch.exp(
                        8 / 9 * 1j * gamma * ((abs(signal) ** 2) + (abs(data_t) ** 2)) * delta_z)

                    # 更改 flag 和 signal 的值
                    data_cpu = data_t.to("cpu").numpy()
                    shared_dict["signal"] = data_cpu
                    shared_dict["flag"] = "y" if shared_dict["flag"] == "x" else "x"


                    # 通知等待的进程
                    condition.notify()  # 通知其他进程条件已满足

                data_fft = fftshift(fft(data_t, dim=0), dim=0) / data_t.size(dim=-1)
                # 1/2 光纤色散
                # 偏振
                data_fft = data_fft * torch.exp(-1j * (beta2 / 2) * (omega ** 2) * (delta_z / 2))
                # 1/2 信号衰减
                data_fft = data_fft * math.exp(-alpha * (delta_z / 2))

            # == EDFA ==
            data_fft = data_fft * math.exp(alpha * span_length)

        tx_signal = ifft(ifftshift(data_fft, dim=0), dim=0) * data_fft.size(dim=-1)
        tx_signal = tx_signal.cpu()
        tx_signal = tx_signal.numpy()
        power_temp = power_temp.cpu()
        power_temp = power_temp.numpy()
        output_queuetype.put((tx_signal, power_temp))



# 在 process_polarization_y 函数中也有类似的逻辑
def fiber_torch(
    signals,
    sampling_rate: float,
    num_spans,
    beta2,
    delta_z,
    gamma,
    alpha,
    span_length,
):
    with multiprocessing.Manager() as manager:
        shared_dict = manager.dict()
        condition = manager.Condition()
        # 初始化 flag 为 "y"
        shared_dict["flag"] = "y"
        shared_dict["signal"] = None
        data_x = signals[0]
        data_y = signals[1]
        # 设置启动方法为'spawn'，确保子进程可以继承主进程的CUDA设备设置
        if get_start_method(allow_none=True) is None:
            set_start_method('spawn')
        result_queuex = Queue()
        result_queuey = Queue()


        # 对 tx_signal_x 处理
        process_x = Process(target=process_polarization, args=(data_x,sampling_rate,num_spans,beta2,delta_z,gamma,
                                                               alpha,span_length,"x",shared_dict,condition,result_queuex,0))

        # process_x.join()


        # 对 tx_signal_y 处理
        process_y = Process(target=process_polarization, args=(data_y, sampling_rate, num_spans, beta2, delta_z, gamma,
                                                               alpha, span_length, "y", shared_dict, condition,
                                                               result_queuey,1))


        process_y.start()
        time.sleep(1)
        process_x.start()
        # process_y.join()
        tx_signal_x, power_temp_x = result_queuex.get()
        tx_signal_y, power_temp_y = result_queuey.get()
        return [tx_signal_x, tx_signal_y], [power_temp_x, power_temp_y]

if __name__ == '__main__':
    #调用in_fiber()函数返回fiber所需的输入参数
    signals,sampling_rate,num_spans,beta2,delta_z,gamma,alpha,span_length=in_fiber()
    #测试fiber函数
    start_time = time.time()
    signals1,signals_power1=fiber_torch(signals,sampling_rate,num_spans,beta2,delta_z,gamma,alpha,span_length)
    end_time = time.time()
    execution_time = end_time - start_time
    with open('results.txt', 'w') as file:
        file.write("signals1:\n")
        for signal in signals1:
            file.write(str(signal) + '\n')

        file.write("signals_power1:\n")
        for power in signals_power1:
            file.write(str(power) + '\n')
    print(f"计算完毕,用时{execution_time}秒")

